import numpy as np


class Meter(object):
    def __init__(self, limit=None):
        self.sum = []
        self.avg = 0
        self.t_num = 0
        self.f_num = 0
        self.limit = limit

    def update(self, v):
        if self.limit:
            assert isinstance(self.limit, int)
            if len(self.sum) == self.limit:
                self.sum.pop(0)
        self.sum.append(v)
        if v == 1:
            self.t_num += 1
        elif v == 0:
            self.f_num += 1
        self.avg = sum(self.sum) / len(self.sum)

    def get_max_index(self):
        if self.sum.count(max(self.sum)) > 1:
            print('more than 1 max value')
        return self.sum.index(max(self.sum))


def get_predict_metrices(sizes, input_size, predict_bboxes, bbox_labels, predict_clses, cls_labels, addition_data=None):
    """

    :param input_size: numpy like [224, 224]
    :param sizes: list of origin image size tuples, like [(224, 224), (512, 512)]
    :param predict_bboxes: numpy with shape [n, 4] representing predict bbox: x1, y1, x2, y2
    :param bbox_labels: numpy with shape [n, 4] representing gt bbox: x1, y1, x2, y2
    :param predict_clses: numpy, shape [n, class_num]
    :param cls_labels: numpy, shape [n, 1]
    :param addition_data: dict of other data
    :return: predict_bbox(numpy shape of [n, 4]),
             gt_bbox(numpy shape of [n, 4]),
             iou(numpy like [0.4, 0.65, ..., 0.5]),
             loc_acc: known ground truth loc-acc,
             loc_map(numpy shape of [n, h_origin, w_origin]),
             top1_loc: ...
             top5_loc: ...
             top1_cls: ...
             top5_cls: ...
    """
    iou, loc_map = [], []
    # detach with label numpy
    bbox_labels = bbox_labels.copy()
    for i, size in enumerate(sizes):
        ''' Calculate IOU '''
        gt_area = np.zeros(input_size)
        x1, y1, x2, y2 = bbox_labels[i, 0], bbox_labels[i, 1], bbox_labels[i, 2], bbox_labels[i, 3]
        gt_area[y1:y2 + 1, x1:x2 + 1] = 1.

        predict_area = np.zeros(input_size)
        x1, y1, x2, y2 = predict_bboxes[i, 0], predict_bboxes[i, 1], predict_bboxes[i, 2], predict_bboxes[i, 3]
        predict_area[y1:y2 + 1, x1:x2 + 1] = 1.

        overlap = (predict_area * gt_area).sum()
        iou.append(overlap / (predict_area.sum() + gt_area.sum() - overlap))

        ''' Resize the bbox to original size '''
        predict_bboxes[i, 1], predict_bboxes[i, 0], predict_bboxes[i, 3], predict_bboxes[i, 2] = \
            int(predict_bboxes[i, 1] * size[0].item() / input_size[0]), int(predict_bboxes[i, 0] * size[1].item() / input_size[1]), \
            int(predict_bboxes[i, 3] * size[0].item() / input_size[0]), int(predict_bboxes[i, 2] * size[1].item() / input_size[1])

        bbox_labels[i, 1], bbox_labels[i, 0], bbox_labels[i, 3], bbox_labels[i, 2] = \
            int(bbox_labels[i, 1] * size[0].item() / input_size[0]), int(bbox_labels[i, 0] * size[1].item() / input_size[1]), \
            int(bbox_labels[i, 3] * size[0].item() / input_size[0]), int(bbox_labels[i, 2] * size[1].item() / input_size[1])
    ''' Classification & Localization Accuracy '''
    # _, predict_index = predict_clses.sort(1, descending=True)
    predict_index = np.flip(np.argsort(predict_clses, axis=-1), axis=-1)

    top1_loc = ((predict_index[:, :1] == cls_labels.reshape(cls_labels.shape[0], 1)).sum(axis=1)
                * (np.array(iou) >= 0.5)).sum().item()
    # top5_loc = ((predict_index[:, :5] == cls_labels.reshape(cls_labels.shape[0], 1)).sum(axis=1)
    #             * (np.array(iou) >= 0.5)).sum().item()

    top1_cls = (predict_index[:, :1] == cls_labels.reshape(cls_labels.shape[0], 1)).sum().item()
    # top5_cls = (predict_index[:, :5] == cls_labels.reshape(cls_labels.shape[0], 1)).sum().item()

    metric_dict = {"predict_bbox": predict_bboxes,
                   "gt_bbox": bbox_labels,
                   "iou": np.array(iou),
                   "predict_cls": predict_clses,
                   "gt_cls": cls_labels,
                   "size": sizes,

                   "gt_loc_acc": np.sum(np.array(iou) >= 0.5),
                   "top1_loc": top1_loc,
                   # "top5_loc": top5_loc,
                   "top1_cls": top1_cls,
                   # "top5_cls": top5_cls,
                   }
    if addition_data:
        metric_dict.update(addition_data)
    return metric_dict


def bbox_predict(heat_map, stride, threshold=0.04):
    """
    Output bbox prediction in model
    """
    max_v = heat_map.max()
    min_v = heat_map.min()
    heat_map = (heat_map - min_v) / (max_v - min_v + 1e-8)
    binary_map = (heat_map >= threshold)
    # binary_map = self.max_connect_area(binary_map)

    bboxes = np.zeros(shape=(heat_map.shape[0], 4), dtype=np.uint16)
    for i, map in enumerate(binary_map):
        ys, xs = np.where(map == 1)
        if len(ys) == 0:
            y1, x1, y2, x2 = 0, 0, 0, 0
            print('No bbox can be found')
        else:
            y1, x1, y2, x2 = ys.min(), xs.min(), ys.max(), xs.max()
        bboxes[i, 0] = int(x1)
        bboxes[i, 1] = int(y1)
        bboxes[i, 2] = int(x2)
        bboxes[i, 3] = int(y2)
    bboxes = bboxes * stride

    return bboxes


def max_connect_area(binary_map):
    m_size = binary_map.shape

    for n in range(m_size[0]):
        f = [i for i in range(m_size[1] * m_size[2])]
        count = np.zeros(shape=[m_size[1] * m_size[2]], dtype=np.int32)

        def find(i):
            if f[i] == i:
                return i
            return find(f[i])

        for i in range(m_size[1]):
            for j in range(m_size[2]):
                if binary_map[n, i, j]:
                    if i > 0 and binary_map[n, i-1, j]:
                        ind = m_size[2] * i + j
                        ind_i_1 = m_size[2] * (i - 1) + j
                        f_c = find(ind_i_1)
                        f[ind] = f_c
                        count[f_c] += 1
                    if j > 0 and binary_map[n, i, j-1]:
                        ind = m_size[2] * i + j
                        ind_j_1 = m_size[2] * i + j - 1
                        f_c = find(ind_j_1)
                        # if the top point is not true
                        if f[ind] == ind:
                            f[ind] = f_c
                            count[f_c] += 1
                        elif f_c != f[ind]:
                            f[f_c] = f[ind]
                            count[f[ind]] += count[f_c]
                            count[f_c] = 0
        max_arg = count.argmax()
        for i in range(m_size[1]):
            for j in range(m_size[2]):
                ind = m_size[2] * i + j
                f_c = find(ind)
                if f_c != max_arg:
                    binary_map[n, i, j] = False
    return binary_map
